import sys
import socket
import getopt
import threading
import subprocess
import time

listen = False
command = False
upload = False
execute = ""
target = ""
upload_destination = ""
port = 0


# 提示的话
def usage():
    print("BHP net Tool")
    print()
    print("Usage: bhpnet.py -t target_host -p port")
    print("-l --listen  \t\t\t- listen on [host]:[port] for incoming connetions")
    print("-e --execute=file_to_run \t- execute the given file upon receiving a connetion")
    print("-c --command \t\t\t- initialize a command shell")
    print("-u --upload=destination  \t- upon receiving connection upload a file and write to [destination]")
    print()
    print()
    print("Examples:")
    print("bnpnet.py -t 192.168.0.1 -p 5555 -l -c")
    print("bnpnet.py -t 192.168.0.1 -p 5555 -l -u=c:\\target.exe")
    print("bnpnet.py -t 192.168.0.1 -p 5555 -l -e=\"cat /etc/passwd\"")
    print("echo 'ABCD' | ./bhpnet.py -t 192.168.11.12 -p 123")
    sys.exit(0)


# 主函数，接受参数，并判断是否监听
def main():
    global listen
    global command
    global upload
    global execute
    global target
    global upload_destination
    global port

    args = ""
    opts = []

    # print(sys.argv)
    if not len(sys.argv[1:]):
        usage()

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hle:t:p:cu", ["help", "listen", "execute", "target", "port",
                                                                "command", "upload"])
    except getopt.GetoptError as err:
        print(str(err))
        usage()

    # print("opts is:", str(opts))
    # print("args is:", str(args))
    # print("args is:", str(sys.argv))

    for o, a in opts:
        if o in ('-h', '--help'):
            usage()
        elif o in ('-l', '--listen'):
            listen = True
        elif o in ('-e', '--execute'):
            execute = a
        elif o in ('-c', '--command'):
            command = True
        elif o in ('-u', '--upload'):
            upload_destination = a
        elif o in ('-t', '--target'):
            target = a
        elif o in ('-p', '--port'):
            port = int(a)
        else:
            assert False, "Unhandled Option"

    if not listen and len(target) > 0 and port > 0:
        # buffer = sys.stdin.read()
        print(args)
        if len(args) is not 0:
            buffer = args[0]
        else:
            buffer = ""
        client_sender(buffer)

    if listen:
        server_loop()


def client_sender(buffer):
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    try:
        client.connect((target, port))

        if len(buffer):
            print("send buffer:" + buffer)
            print(b"send buffer:" + bytes(buffer, 'utf-8'))
            print('run send.')
            if not buffer.endswith('\n'):
                buffer += '\n'
            client.send(bytes(buffer, 'utf-8'))

        while True:
            recv_len = 4096
            response = ""

            # time.sleep(1)
            while recv_len >= 4096 or response == "":
                data = client.recv(4096)
                recv_len = len(data)
                response += str(data, encoding='utf-8')

            print(response, end="")

            buffer = input("")
            if not buffer.endswith('\n'):
                buffer += '\n'
            print(bytes(buffer, 'utf-8'))
            client.send(bytes(buffer, 'utf-8'))
    except Exception as e:
        print(e)
        print("[*] Exception! Exiting.")


def server_loop():
    print("run server.")
    global target
    if not target:
        target = "0.0.0.0"

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((target, port))
    server.listen(5)
    print(target)
    print(port)

    while True:
        client_socket, addr = server.accept()
        print(1)

        client_thread = threading.Thread(target=client_handler, args=(client_socket, ))
        client_thread.start()
    pass


def run_command(command):
    command = command.strip()
    try:
        output = subprocess.check_output(command, stderr=subprocess.STDOUT, shell=True)
        print('output is: {};  type is:{}'.format(output, type(output)))
    except Exception as e:
        print(e)
        output = "Failed to execute command.\r\n"
    return output


def client_handler(client_socket):
    global upload
    global execute
    global command

    if len(upload_destination):
        file_buffer = ""

        while True:
            data = client_socket.recv(1024)

            if not data:
                break
            else:
                file_buffer += str(data)

        try:
            file_descriptor = open(upload_destination, "wb")
            file_descriptor.write(file_buffer)
            file_descriptor.close()

            client_socket.send("Successfully saved file to %s\r\n" % upload_destination)
        except:
            client_socket.send("Failed to save file to %s\r\n" % upload_destination)

    if len(execute):
        output = run_command(execute)
        client_socket.send(bytes(output, 'utf-8'))

    if command:
        print("run command.")
        client_socket.send(b"<BHP:#> ")

        while True:
            tmp_msg = ""
            cmd_buffer = ""
            while '\n' not in cmd_buffer:
                # print(cmd_buffer)
                tmp_msg = client_socket.recv(1024)
                # print("buffer is: {}\n type is:{}".format(cmd_buffer, type(cmd_buffer)))
                print("client_socket.recv is: {}\n type is:{}".format(tmp_msg, type(tmp_msg)))
                cmd_buffer += str(tmp_msg, encoding='utf-8')

                print('tmp:' + str(tmp_msg))

                if tmp_msg == b"":
                    cmd_buffer += b'\n'
                    break
                elif tmp_msg == b'exit\n':
                    client_socket.send(b'server is exited.')
                    client_socket.close()
                    break
            print('tmp:' + str(tmp_msg))

            if tmp_msg == b'exit\n':
                break

            print("final is:" + cmd_buffer)

            response = run_command(cmd_buffer)
            if isinstance(response, str):
                response = bytes(response, encoding='utf-8')
            client_socket.send(response + b"\n<BHP:#> ")


if __name__ == '__main__':
    main()
